import json

from flask import Flask, render_template, request, jsonify, send_file
import os
import time
import numpy as np
import cv2
from PIL import Image
import io
import base64
from agent import ImageFusionAgent

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key-here'
app.config['UPLOAD_FOLDER'] = 'temp_uploads'

# 确保上传目录存在
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

# 初始化AI Agent
try:
    from agent import fuse_color_images
    from joint import ImageStitcher
    import panorama

    def adapted_fusion_function(near_img_path, far_img_path):
        success, fused_img = fuse_color_images(near_img_path, far_img_path)
        return success, fused_img

    def stitch_function(img_paths):
        # imgs = []
        # for path in img_paths:
        #     img = cv2.imread(path)
        #     imgs.append(img)
        # stitch = ImageStitcher('panorama')
        # success, result = stitch.stitch_images(imgs)
        success, result = panorama.stitch(img_paths)
        print(success)
        return success, result

    agent = ImageFusionAgent(adapted_fusion_function, stitch_function)
    print("AI Agent初始化成功")
except Exception as e:
    print(f"初始化AI Agent失败: {str(e)}")
    agent = None

# 存储用户会话状态
user_sessions = {}

def get_user_session(session_id):
    """获取或创建用户会话"""
    if session_id not in user_sessions:
        user_sessions[session_id] = {
            'conversation_history': [],
            # 'image_paths': [],  # 改为图像路径列表
            'result_image': None,
            'last_activity': time.time()
        }
    return user_sessions[session_id]


def cleanup_old_sessions():
    """清理过期会话"""
    current_time = time.time()
    expired_sessions = []
    for session_id, session_data in user_sessions.items():
        if current_time - session_data['last_activity'] > 3600:  # 1小时过期
            expired_sessions.append(session_id)

    for session_id in expired_sessions:
        del user_sessions[session_id]


def image_to_base64(image):
    """将图像转换为base64字符串"""
    if image is None:
        return None

    try:
        if isinstance(image, np.ndarray):
            # 如果是OpenCV格式(BGR)，转换为RGB
            if len(image.shape) == 3 and image.shape[2] == 3:
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(image)

        # 转换为base64
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return f"data:image/png;base64,{img_str}"
    except Exception as e:
        print(f"图像转换错误: {str(e)}")
        return None


@app.route('/')
def index():
    """主页面"""
    return render_template('index.html')


@app.route('/upload_images', methods=['POST'])
def upload_images():
    """上传图像"""
    if agent is None:
        return jsonify({'success': False, 'error': 'AI Agent未初始化'})

    try:
        session_id = request.form.get('session_id', 'default')
        files = request.files.getlist('images')

        if not files or len(files) == 0:
            return jsonify({'success': False, 'error': '缺少图像文件'})

        # 获取用户会话（仅用于活动跟踪）
        session = get_user_session(session_id)
        session['last_activity'] = time.time()

        previews = []
        paths = []
        for i, file in enumerate(files):
            # 保存图像文件
            filename = f"image_{i}_{int(time.time())}_{session_id}.png"
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            image = Image.open(file.stream)
            image.save(filepath)

            # 生成预览
            image.thumbnail((150, 150))
            preview_base64 = image_to_base64(image)

            previews.append(preview_base64)
            paths.append(filepath)

        return jsonify({
            'success': True,
            'previews': previews,
            'paths': paths,
            'count': len(files)
        })

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})


@app.route('/chat', methods=['POST'])
def chat():
    """处理对话"""
    if agent is None:
        return jsonify({'success': False, 'error': 'AI Agent未初始化'})

    try:
        session_id = request.form.get('session_id', 'default')
        user_input = request.form.get('message')
        image_paths = request.form.getlist('image_paths')

        if not user_input:
            return jsonify({'success': False, 'error': '消息不能为空'})

        # 获取用户会话
        session = get_user_session(session_id)
        session['last_activity'] = time.time()

        # 构建完整消息（包含所有图像路径）
        full_message = user_input
        # if 'image_paths' in session and session['image_paths']:
        #     full_message += f" 图像路径: {', '.join(session['image_paths'])}"
        if image_paths:
            full_message += f" 图像路径: {', '.join(image_paths)}"

        # 处理AI响应
        response = agent.chat(full_message)

        # 更新对话历史
        session['conversation_history'].append({
            'role': 'user',
            'content': user_input,
            'time': time.strftime('%H:%M:%S')
        })

        session['conversation_history'].append({
            'role': 'assistant',
            'content': response.get('response', '无响应'),
            'time': time.strftime('%H:%M:%S')
        })

        # 检查是否需要执行工具
        tool_info = None
        if response.get("tool") and response["tool"] != "none":
            tool_info = {
                'tool': response['tool'],
                'tool_input': response.get('tool_input', {}),
                'image_paths': image_paths
            }

        return jsonify({
            'success': True,
            'response': response.get('response', '无响应'),
            'tool_info': tool_info,  # 告诉前端需要执行什么工具
            'conversation_history': session['conversation_history'][-10:]
        })

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

@app.route('/clear_conversation', methods=['POST'])
def clear_conversation():
    """清空对话"""
    try:
        session_id = request.form.get('session_id', 'default')

        # 获取用户会话
        session = get_user_session(session_id)
        session['last_activity'] = time.time()

        # 清空对话历史
        session['conversation_history'] = []
        session['result_image'] = None

        return jsonify({'success': True})

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})


@app.route('/delete_image', methods=['POST'])
def delete_image():
    """删除单个图像文件"""
    try:
        session_id = request.form.get('session_id', 'default')
        image_path = request.form.get('image_path')

        if not image_path:
            return jsonify({'success': False, 'error': '缺少图像路径'})

        # 安全检查：确保图像路径在允许的目录内
        if not image_path.startswith(app.config['UPLOAD_FOLDER']):
            return jsonify({'success': False, 'error': '无权删除该文件'})

        if os.path.exists(image_path):
            os.remove(image_path)
            return jsonify({'success': True})
        else:
            return jsonify({'success': False, 'error': '文件不存在'})

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})


@app.route('/delete_images', methods=['POST'])
def delete_images():
    """删除多个图像文件"""
    try:
        session_id = request.form.get('session_id', 'default')
        image_paths = request.form.getlist('image_paths')

        deleted_count = 0
        errors = []

        for image_path in image_paths:
            # 安全检查：确保图像路径在允许的目录内
            if not image_path.startswith(app.config['UPLOAD_FOLDER']):
                errors.append(f"无权删除文件: {image_path}")
                continue

            if os.path.exists(image_path):
                os.remove(image_path)
                deleted_count += 1

        return jsonify({
            'success': True,
            'deleted_count': deleted_count,
            'errors': errors
        })

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})


@app.route('/execute_tool', methods=['POST'])
def execute_tool():
    """异步执行工具并返回结果图像"""
    if agent is None:
        return jsonify({'success': False, 'error': 'AI Agent未初始化'})

    try:
        session_id = request.form.get('session_id', 'default')
        tool = request.form.get('tool')
        tool_input_json = request.form.get('tool_input')
        image_paths = request.form.getlist('image_paths')

        # 获取用户会话
        session = get_user_session(session_id)
        session['last_activity'] = time.time()

        # 解析工具输入
        tool_input = json.loads(tool_input_json) if tool_input_json else {}
        # print(tool_input)
        # 执行工具
        tool_result = agent._execute_tool({
            'tool': tool,
            'tool_input': tool_input
        })

        # 检查处理结果
        result_image_base64 = None
        if tool == 'fuse_images':
            if hasattr(agent, 'last_fusion_result') and agent.last_fusion_result:
                success, result_image = agent.last_fusion_result
                if success and result_image is not None:
                    session['result_image'] = result_image
                    result_image_base64 = image_to_base64(result_image)
        elif tool == 'stitch_images':
            if hasattr(agent, 'last_stitch_result') and agent.last_stitch_result:
                success, result_image = agent.last_stitch_result
                if success and result_image is not None:
                    session['result_image'] = result_image
                    result_image_base64 = image_to_base64(result_image)

        return jsonify({
            'success': True,
            'tool_result': tool_result,
            'result_image': result_image_base64,
            'tool_used': tool
        })

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

@app.route('/download_result', methods=['POST'])
def download_result():
    """下载融合结果"""
    try:
        session_id = request.form.get('session_id', 'default')

        # 获取用户会话
        session = get_user_session(session_id)
        session['last_activity'] = time.time()

        if session['result_image'] is None:
            return jsonify({'success': False, 'error': '没有可下载的融合结果'})

        # 将图像保存为字节流
        fused_image = session['result_image']
        if isinstance(fused_image, np.ndarray):
            # 如果是OpenCV格式，转换为PIL图像
            if len(fused_image.shape) == 3 and fused_image.shape[2] == 3:
                fused_image = cv2.cvtColor(fused_image, cv2.COLOR_BGR2RGB)
            fused_image = Image.fromarray(fused_image)

        # 保存到内存
        img_io = io.BytesIO()
        fused_image.save(img_io, 'PNG')
        img_io.seek(0)

        # 返回文件
        return send_file(
            img_io,
            mimetype='image/png',
            as_attachment=True,
            download_name='result_img.png'
        )

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)